{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# TiDB\n",
    "\n",
    "> [TiDB](https://github.com/pingcap/tidb) is an open-source, cloud-native, distributed, MySQL-Compatible database for elastic scale and real-time analytics.\n",
    "\n",
    "This notebook introduces how to use TiDB to store chat message history. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "Firstly, we will install the following dependencies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --upgrade --quiet langchain langchain_openai"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Configuring your OpenAI Key"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import getpass\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"Input your OpenAI API key:\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we will configure the connection to a TiDB. In this notebook, we will follow the standard connection method provided by TiDB Cloud to establish a secure and efficient database connection."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# copy from tidb cloud console\n",
    "tidb_connection_string_template = \"mysql+pymysql://<USER>:<PASSWORD>@<HOST>:4000/<DB>?ssl_ca=/etc/ssl/cert.pem&ssl_verify_cert=true&ssl_verify_identity=true\"\n",
    "tidb_password = getpass.getpass(\"Input your TiDB password:\")\n",
    "tidb_connection_string = tidb_connection_string_template.replace(\n",
    "    \"<PASSWORD>\", tidb_password\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating historical data\n",
    "\n",
    "Creating a set of historical data, which will serve as the foundation for our upcoming demonstrations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime\n",
    "\n",
    "from langchain_community.chat_message_histories import TiDBChatMessageHistory\n",
    "\n",
    "history = TiDBChatMessageHistory(\n",
    "    connection_string=tidb_connection_string,\n",
    "    session_id=\"code_gen\",\n",
    "    earliest_time=datetime.utcnow(),  # Optional to set earliest_time to load messages after this time point.\n",
    ")\n",
    "\n",
    "history.add_user_message(\"How's our feature going?\")\n",
    "history.add_ai_message(\n",
    "    \"It's going well. We are working on testing now. It will be released in Feb.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content=\"How's our feature going?\"),\n",
       " AIMessage(content=\"It's going well. We are working on testing now. It will be released in Feb.\")]"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "history.messages"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Chatting with historical data\n",
    "\n",
    "Let’s build upon the historical data generated earlier to create a dynamic chat interaction.  \n",
    "\n",
    "Firstly, Creating a Chat Chain with LangChain:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder\n",
    "from langchain_openai import ChatOpenAI\n",
    "\n",
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\n",
    "            \"system\",\n",
    "            \"You're an assistant who's good at coding. You're helping a startup build\",\n",
    "        ),\n",
    "        MessagesPlaceholder(variable_name=\"history\"),\n",
    "        (\"human\", \"{question}\"),\n",
    "    ]\n",
    ")\n",
    "chain = prompt | ChatOpenAI()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Building a Runnable on History:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.runnables.history import RunnableWithMessageHistory\n",
    "\n",
    "chain_with_history = RunnableWithMessageHistory(\n",
    "    chain,\n",
    "    lambda session_id: TiDBChatMessageHistory(\n",
    "        session_id=session_id, connection_string=tidb_connection_string\n",
    "    ),\n",
    "    input_messages_key=\"question\",\n",
    "    history_messages_key=\"history\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Initiating the Chat:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AIMessage(content='There are 31 days in January, so there are 30 days until our feature is released in February.')"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "response = chain_with_history.invoke(\n",
    "    {\"question\": \"Today is Jan 1st. How many days until our feature is released?\"},\n",
    "    config={\"configurable\": {\"session_id\": \"code_gen\"}},\n",
    ")\n",
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checking the history data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[HumanMessage(content=\"How's our feature going?\"),\n",
       " AIMessage(content=\"It's going well. We are working on testing now. It will be released in Feb.\"),\n",
       " HumanMessage(content='Today is Jan 1st. How many days until our feature is released?'),\n",
       " AIMessage(content='There are 31 days in January, so there are 30 days until our feature is released in February.')]"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "history.reload_cache()\n",
    "history.messages"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "langchain",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
